package org.apache.zookeeper;

import org.apache.jute.BinaryInputArchive;
import org.apache.zookeeper.AsyncCallback.GetDataCallback;
import org.apache.zookeeper.AsyncCallback.SetDataCallback;
import org.apache.zookeeper.ZooDefs.OpCode;
import org.apache.zookeeper.proto.*;
import org.apache.zookeeper.server.paxos2.Configuration;
import org.apache.zookeeper.server.paxos2.Mapping;

import java.io.IOException;
import java.util.*;

public class BlockStorage {

    ClientCnxn servers;

    private final int outstandingLimit = Configuration.clientOutStanding;

    private int clientId;
    private int writeCount;

    private class PendingWrite extends TimeoutHandler {
        public long blockId;
        public byte[] data;

        public int xid;
        public int phase;
        public HashSet<Integer> prepareReceived = new HashSet<Integer>();
        public HashSet<Integer> writeReceived = new HashSet<Integer>();

        public SetDataCallback cb;
        public Object ctx;
        public boolean writeSent = false;

        public PendingWrite(long blockId, int xid, byte[] data, SetDataCallback cb, Object ctx) {
            super(Configuration.timeoutInterval);
            this.blockId = blockId;
            this.xid = xid;
            this.data = data;

            this.cb = cb;
            this.ctx = ctx;
        }

        public synchronized void addPrepareReply(int serverId) {
            prepareReceived.add(serverId);
        }

        public synchronized boolean isDataReady() {
            return prepareReceived.size() >= Configuration.writeQuorumSize;
        }

        public synchronized void addWriteReply(int serverId) {
            writeReceived.add(serverId);
        }

        public synchronized boolean isWriteReady() {
            return writeReceived.size() >= Configuration.writeQuorumSize;
        }

        public synchronized void handleTimeout() {
            //if (true) return;
            System.out.println("timeout write client=" + clientId + " xid=" + xid + " dataReady="+isDataReady()+" writeSent="+writeSent);
            LinkedList<Integer> liveServers = servers.listLiveServers();
            if (!isDataReady()) {
                PrepareDataRequest prepare = new PrepareDataRequest(clientId, 2, data);
                RequestHeader header = new RequestHeader(xid, xid, OpCode.prepareData);
                for (Integer node : liveServers) {
                    //System.out.println("send prepare "+xid+" "+System.currentTimeMillis());
                    servers.queuePacket(header, prepare, node);
                }
                this.resetTimeout(Configuration.timeoutInterval);
            } else {
                if (writeSent && liveServers.size() > 0) {
                    SetDataRequest set = new SetDataRequest(blockId, new byte[1]);
                    RequestHeader h = new RequestHeader(xid, xid, OpCode.setData);
		    for(Integer sid : liveServers){
			//System.out.println("resent to "+sid);
                    	servers.queuePacket(h, set, sid);
		    }
                }
                this.resetTimeout(Configuration.timeoutInterval);
            }
        }

    }

    private class PendingRead extends TimeoutHandler {
        public byte[] data;
        public long blockId;

        public int xid;
        public long startTime;
        public HashSet<Integer> nodeReceived = new HashSet<Integer>();

        public GetDataCallback cb;
        public Object ctx;

        public PendingRead(long blockId, int xid, GetDataCallback cb, Object ctx) {
            super(Configuration.timeoutInterval);
            this.blockId = blockId;
            this.xid = xid;
            this.startTime = System.currentTimeMillis();

            this.cb = cb;
            this.ctx = ctx;
        }

        public void addReply(int serverId) {
            nodeReceived.add(serverId);
        }

        public boolean isDataReady() {
            return nodeReceived.size() >= 1;//Configuration.writeQuorumSize;
        }

        public synchronized void handleTimeout() {
            System.out.println("timeout read client=" + clientId + " xid=" + xid);
            BoloskyGetDataRequest request = new BoloskyGetDataRequest(blockId, 0); //need fix
            RequestHeader header = new RequestHeader(xid, xid, OpCode.boloskyGetData);
            //List<Integer> quorum = Mapping.keyToNode(blockId);
            servers.queuePacket(header, request, 3);
            this.resetTimeout(Configuration.timeoutInterval);

        }

    }

    private final TreeMap<Integer, PendingWrite> pendingWrites = new TreeMap<Integer, PendingWrite>();

    private final TreeMap<Integer, PendingRead> pendingReads = new TreeMap<Integer, PendingRead>();

    private final TimeoutChecker timeoutChecker = new TimeoutChecker();

    public BlockStorage(String hosts, int id) throws IOException {
        this.clientId = id;
        servers = new ClientCnxn(hosts, id);
        servers.start(this);
        try {
            Thread.sleep(3000);
        } catch (Exception e) {
        }

        /*SetDataRequest set = new SetDataRequest(-1, new byte[1]);
        RequestHeader h = new RequestHeader(getXid(), OpCode.setData);
        servers.queuePacket(h, set, 3);*/
        //write(-1, new byte[1]);
        timeoutChecker.start();
    }

    private int xid = 0;

    synchronized private int getXid() {
        return xid++;
    }

    synchronized public void initializeXid(int xid) {
        if (xid > this.xid)
            this.xid = xid;
    }


    public void write(long blockId, byte[] data, SetDataCallback cb, Object ctx) throws KeeperException,
            InterruptedException {
        int xid;
        synchronized (pendingWrites) {
            while (pendingWrites.size() >= outstandingLimit) {
                try {
                    pendingWrites.wait();
                }
                catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
            xid = getXid();
            PendingWrite p = new PendingWrite(blockId, xid, data, cb, ctx);
            pendingWrites.put(xid, p);
            timeoutChecker.addEvent(p);
        }
        PrepareDataRequest prepare = new PrepareDataRequest(clientId, 0, data);
        RequestHeader header = new RequestHeader(xid, xid, OpCode.prepareData);
        List<Integer> writeQuorum = Mapping.keyToNode(blockId, servers.listLiveServers());
        for (Integer node : writeQuorum) {
            //System.out.println("send prepare to "+node);
            servers.queuePacket(header, prepare, node);
        }
    }

    long count = 0;


    public void read(long blockId, GetDataCallback cb, Object ctx) throws KeeperException,
            InterruptedException {
        System.out.println("read " + blockId);
        int xid;
        synchronized (pendingReads) {
            while (pendingReads.size() >= outstandingLimit) {
                try {
                    pendingReads.wait();
                }
                catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
            xid = getXid();
            PendingRead p = new PendingRead(blockId, xid, cb, ctx);
            pendingReads.put(xid, p);
            timeoutChecker.addEvent(p);
        }
        BoloskyGetDataRequest request = new BoloskyGetDataRequest(blockId, -1); //need fix
        RequestHeader header = new RequestHeader(xid, xid, OpCode.boloskyGetData);
        //List<Integer> quorum = Mapping.keyToNode(blockId);
        servers.queuePacket(header, request, 3);
        /*
           * int id = Mapping.keyToNode(blockId).get(0) - 1; return
           * zks[id].getData(String.valueOf(blockId), false, new Stat());
           */
    }

    public void handleResponse(ReplyHeader header, BinaryInputArchive bbia, int serverId) throws IOException {
        //System.out.println("Got "+header.getXid()+" "+header.getType()+" from"+serverId);
        switch (header.getType()) {
            case OpCode.prepareData:
                PendingWrite p = null;
                synchronized (pendingWrites) {
                    p = pendingWrites.get(header.getXid());
                }
                if (p == null) {
                    //System.out.println("prepareData "+header.getXid()+" does not exist");
                    break;
                }
                if (!p.isDataReady()) {
                    p.addPrepareReply(serverId);
                    if (p.isDataReady()) {
                        //System.out.println("Prepare ready "+header.getXid());
                        /*SetDataRequest set = new SetDataRequest(p.blockId, new byte[1]);
                      RequestHeader h = new RequestHeader(p.xid, OpCode.setData);
                      servers.queuePacket(h, set, 3);*/
                        synchronized (pendingWrites) {
                            for (Integer xid : pendingWrites.keySet()) {
                                PendingWrite tmp = pendingWrites.get(xid);
                                if (tmp.isDataReady()) {
				    if(!tmp.writeSent){
                                    SetDataRequest set = new SetDataRequest(tmp.blockId, new byte[1]);
                                    RequestHeader h = new RequestHeader(tmp.xid, tmp.xid, OpCode.setData);
                                    //System.out.println(clientId + " sends setdata " + xid + " " + System.currentTimeMillis());
                                    servers.queuePacket(h, set, 3);
                                    tmp.writeSent = true;
				    }
                                } else{
				    //System.out.println(clientId + " "+xid+" is not ready");
                                    break;
				}
                            }

                        }
                    }
                }
                break;
            case OpCode.setData:
                p = null;
                synchronized (pendingWrites) {
                    p = pendingWrites.get(header.getXid());
                    if (p != null) {
                        p.addWriteReply(serverId);
                        if (p.isWriteReady()) {
                            pendingWrites.remove(header.getXid());
                            p.cb.processResult(header.getErr(), p.blockId, p.ctx);
                            p.finish();
                            pendingWrites.notify();
                        }
                    }
                }

                //System.out.println("setData "+header.getXid()+" finished with code "+header.getErr());

                break;
            case OpCode.getData:
                GetDataResponse response = new GetDataResponse();
                response.deserialize(bbia, "response");
                System.out.println("getdata xid=" + header.getXid());
                PendingRead pRead = null;
                synchronized (pendingReads) {
                    pRead = pendingReads.get(header.getXid());
                }
                if (pRead == null) {
                    break;
                } else
                    pRead.data = response.getData();
                synchronized (pendingReads) {
                    Iterator<Map.Entry<Integer, PendingRead>> iter = pendingReads.entrySet().iterator();
                    while (iter.hasNext()) {
                        PendingRead tmp = iter.next().getValue();
                        if (tmp.data != null) {
                            tmp.cb.processResult(header.getErr(), tmp.blockId, tmp.ctx, response.getData());
                            iter.remove();
                        } else
                            break;
                    }
                    pendingReads.notify();
                }
                break;
            case OpCode.boloskyGetData:
                BoloskyGetDataResponse resp = new BoloskyGetDataResponse();
                resp.deserialize(bbia, "response");
                pRead = null;
                synchronized (pendingReads) {
                    pRead = pendingReads.get(header.getXid());
                }
                if (pRead == null) {
                    System.out.println("Redundant bolosky read reply");
                    break;
                } else
                    pRead.data = resp.getData();
                synchronized (pendingReads) {
                    Iterator<Map.Entry<Integer, PendingRead>> iter = pendingReads.entrySet().iterator();
                    while (iter.hasNext()) {
                        PendingRead tmp = iter.next().getValue();
                        if (tmp.data != null) {
                            tmp.cb.processResult(header.getErr(), tmp.blockId, tmp.ctx, resp.getData());
                            tmp.finish();
                            iter.remove();
                        } else
                            break;
                    }
                    pendingReads.notify();
                }
                break;
        }
    }


}
